#!/usr/bin/env python

# -------
# Imports
# -------
import config

# ROS 
import roslib; roslib.load_manifest(config.NODE_NAME)
import rospy
from geometric_image_filter.msg import FilteredImage
from sensor_msgs.msg import Image, CameraInfo, LaserScan, RegionOfInterest, PointCloud2
from geometry_msgs.msg import Point, PointStamped
from cv_bridge import CvBridge, CvBridgeError
import cv
import tf

# General  python
import threading
import sys
import struct
import math
from numpy import *


# ------------------------------------------
class GeometricImageFilter(threading.Thread):
# ------------------------------------------
	# Publishers/Subscriptions
	# ------------------------
	filteredImagePub = None 		# Publisher for the FilteredImage messages this node outputs
	transformListener = None		# This is used to access the tf frames
	bridge = None					# cv_bridge for converting images into a usable format

	# Stored messages & relevant information
	# --------------------------------------
	prevCameraInfoMsg = None 		# The last camera info message that we
	cameraInfoMsg = None			# The most recent camera info message we've received 
	imageMsg = None					# The most recent image message we've received
	stereoPtCloudMsg = None			# The most recent stereo data we've received
	laserScanMsg = None				# The most recent laser data we've received (laser scan)
	cameraPtWorld = None			# The most recent computation of where the camera is physically in the world frame

	# Information about filtered images that we don't want to have to re-alocate/compute
	# ---------------------------------------------------------------------------------
	wRange = None					# A list of the object window sizes
	wvMask = None					# A mapping from (w,v) pairs to a boolean indicating whether or not that region is geometrically feasible
	wvNumRegions = {} 				# A mapping from (w,v) pairs to the number of regions of interest that pair represents (aross all values of u)

	# Back-end filtering variables
	# ----------------------------
	groundNormalWorld = array([[0],[0],[1],[0]])	# The ground plane normal vector (in the world coordinate frame)
	ry = None										# Projection of v into the real world (for each v)
	dS = None										# Depth range based on physical size
	dH = None										# Depth range based on physical height
	dC = None 										# The intersection of dS and dH
	mS = {} 										# Stereo disparity depth constraint
	
	# ----------------
	def __init__(self):
	# ----------------
		# Super
		threading.Thread.__init__(self)

		# Initialize the node
		rospy.init_node(config.NODE_NAME)

		# Subscriptions
		rospy.Subscriber(config.CAMERA_INFO_TOPIC, CameraInfo, self.__cameraInfoCallback)
		rospy.Subscriber(config.IMAGE_TOPIC, Image, self.__imageCallback)
		rospy.Subscriber(config.STEREO_PT_CLOUD_TOPIC, PointCloud2, self.__stereoPtCloudCallback)
		rospy.Subscriber(config.LASER_SCAN_TOPIC, LaserScan, self.__laserScanCallback)
		self.transformListener = tf.TransformListener()

		# Create the cv bridge
		self.bridge = CvBridge()

		# We're going to publish filtered images
		self.filteredImagePub = rospy.Publisher(config.FILTERED_IMAGE_TOPIC, FilteredImage)

	# ----
	# Core|
	# -----------
	def run(self):
	# -----------
		""" The main control loop for the geometric_image_processing node. """

		from time import time

		# Status info
		rospy.loginfo('Starting %s node' % (config.NODE_NAME))

		# Preprocessing
		# -------------
		# Block for a camera info message
		rospy.loginfo('Blocking for camera calibration information...')
		while self.cameraInfoMsg == None: pass
		cameraInfoMsg = self.cameraInfoMsg

		# Determine the minimum object size in *pixels*
		minWindowSize = self.__calculateMinWindowSize(cameraInfoMsg.K, config.MIN_SIZE_M, config.MAX_DEPTH_M)
		rospy.loginfo('Minimum window size: %ix%i pixels', minWindowSize, minWindowSize)

		# Precompute the (v,w) mask
		rospy.loginfo('Generating set of possible regions of interest')
		self.wRange = [w for w in self.__windowSizeRange(minWindowSize, min(cameraInfoMsg.width, cameraInfoMsg.height), config.WINDOW_SCALE_FACTOR)]
		self.wvMask = self.__generateWVMask(self.wRange, self.cameraInfoMsg.width, self.cameraInfoMsg.height, self.wvNumRegions)

		# Main control loop
		# -----------------
		i = 0
		times = []
		averages = []
		rospy.loginfo('Entering main control loop')
		while not rospy.is_shutdown():

			# Collect Messages:
			# ----------------
			imageMsg, cameraInfoMsg, cameraPtWorld, groundNormalCamera, stereoPtCloudMsg, laserScanMsg = self.__getMessages()

			# Make sure we have enough info to proceed
			if imageMsg != None and cameraInfoMsg != None and cameraPtWorld != None:
				###################
				startTime = time() 
				###################

				updateD = False	
		
				# Update size constraint:	
				# ----------------------
				# If necessary, compute the depths consistent with the possible image window sizes and the min/max feasible object size
				if self.prevCameraInfoMsg == None or self.prevCameraInfoMsg.K != cameraInfoMsg.K:
					rospy.loginfo('Computing the range of feasible depths given image window size and feasible object size')
					self.dS = self.computeSizeConstraint(cameraInfoMsg, self.wRange, config.MIN_SIZE_M, config.MAX_SIZE_M)
					self.dS = self.__filterDepthConstraints(self.dS, config.MAX_DEPTH_M)

					# Update the projections of v into the real world 	
					self.ry = self.__computeProjection(cameraInfoMsg)					
	
					# Update the previously processed camera info message
					self.prevCameraInfoMsg = cameraInfoMsg

					# Since we've changed dS, we're going to have to update dC
					updateD = True

				# Update height constraint:
				# ------------------------
				# If necessary, compute the depths consistent with the possible image heights and the min/max real world feasible object height
				if self.cameraPtWorld == None or cameraPtWorld[2] != self.cameraPtWorld[2]:
					rospy.loginfo('Computing the range of feasible depths given image height and feasible object height')
					self.dH = self.computeHeightConstraint(cameraInfoMsg, groundNormalCamera, self.ry, cameraPtWorld[2], config.MIN_HEIGHT_M, config.MAX_HEIGHT_M)
					self.dH = self.__filterDepthConstraints(self.dH, config.MAX_DEPTH_M)

					# Update the camera point
					self.cameraPtWorld = cameraPtWorld
				
					# Since we've changed dH, we're going to have to update dC 
					updateD = True

				# Update the combined depth constraints:
				# -------------------------------------
				if updateD:
					rospy.loginfo('Computing the intersection of the size & height constraints')					
					self.dC = self.combineDepthConstraints(self.dS, self.dH)

					############################
					# Summary image generation #
					#######################################################################################################
					#summaryImage = self.generateSummaryImage(imageMsg, self.wvMask, self.wRange, self.cameraInfoMsg.height)
					#cv.SaveImage('data/%i.png' % i, summaryImage)
					#i += 1
					#######################################################################################################

				# Process the stereo message
				if stereoPtCloudMsg != None:
					self.processStereoMessage(stereoPtCloudMsg,self.dH, self.mS)

				# Filter the set of possible ROI's:
				# --------------------------------
				numFeasibleRegions = self.filterRegionsOfInterest(self.wvMask, self.wRange, self.wvNumRegions, self.dC)

				# Timing information
				duration = time() - startTime
				times.append(duration)
				averages.append(float(numFeasibleRegions) / self.totalROI)

		rospy.loginfo('Shutting down %s node' % (config.NODE_NAME))

		# Status info
		rospy.loginfo('Average duration: %s (%s)', str(sum(times[:-1]) / (len(times) - 1)), str((len(times) - 1)/sum(times[:-1]))) 
		rospy.loginfo('Average percent of total: %f', sum(averages) / len(averages))
		rospy.loginfo('Total number of frames processed: %s', len(times))
		rospy.loginfo('Computing & dispalying summary image...')
		
	# -------------------------------------------------------------------------------------
	def computeSizeConstraint(self, cameraInfoMsg, windowSizes, minWorldSize, maxWorldSize):
	# -------------------------------------------------------------------------------------
		# Get the focal length, fx
		fx = cameraInfoMsg.K[0]		

		# Common values
		sMinFx = fx * minWorldSize
		sMaxFx = fx * maxWorldSize
		
		# Compute the range of feasible depths
		feasibleSizes = {}
		for w in windowSizes:
			val = ((sMinFx / w), (sMaxFx / w))
			feasibleSizes[w] = val		

		return feasibleSizes

	# -------------------------------------------------------------------------------------------------------------
	def computeHeightConstraint(self, cameraInfoMsg, groundNormalCamera, ry, cameraHeightM, minHeightM, maxHeightM):
	# -------------------------------------------------------------------------------------------------------------
		# Common values
		dH = {}
		hMinDiff = (minHeightM - cameraHeightM)
		hMaxDiff = (maxHeightM - cameraHeightM)

		# Alias the ground plane's normal vector
		n = groundNormalCamera

		# Comptute the depth constraint
		for v in range(0, cameraInfoMsg.height):
			nr = (float(n[1]) * ry[v]) + float(n[2])
			dHMin = (hMinDiff / nr)
			dHMax = (hMaxDiff / nr)
			val = (min(dHMin, dHMax), max(dHMin, dHMax)) 

			dH[v] = val

		return dH

	# ---------------------------------------
	def combineDepthConstraints(self, dS, dH):
	# ---------------------------------------
		""" Combines the size and height geometric constraints into a single collection.
		Specifically, this method combines dS and dH in to a mapping of (w,v) --> (min, max depth).
		"""

		# The combination of dS and dH
		dC = {}

		# Combine the constraints
		for w, sizeRange in dS.iteritems():
			for v, heightRange in dH.iteritems():
				dC[w,v] = self.__rangeIntersection(sizeRange, heightRange)

		return dC

	# -----------------------------------------
	def processStereoMessage(self, msg, dH, mS):
	# -----------------------------------------

		for w in range(0, 15):
			for v in range (0, msg.height):
				for u in range(0, msg.width):
					pass

		"""
		fmt = ''.join(['f' for i in range(0, msg.width * msg.height)])
		data = struct.unpack_from(fmt, msg.data)
		dataMask = map(math.isnan, data)
		#total = sum(map(lambda x: 0 if math.isnan(x) else 1, data))	
		rospy.loginfo(total / (640.0 * 480))
		"""

		"""
		# Pull out some things that aren't going to change
		rowStep = msg.row_step
		ptStep = msg.point_step
		data = msg.data
		width = msg.width
		height = msg.height
		fmt = 'f'

		# Methods
		unpack_from = struct.unpack_from
		isnan = math.isnan

		offset = 0
		for v in range(0, height):
			# Only consider the row if there's a valid entry in dH
			if dH[v] != None:

				# Loop over the row in the image
				for u in range(0, width):
					# Unpack the depth to a float
					dTup = unpack_from(fmt, data, offset)
					d = dTup[0]
					offset += ptStep

					# Only process further if we have a valid depth
					if not isnan(d):
						mS[v,u] = d
		"""

	# -----------------------------------------------------------------
	def filterRegionsOfInterest(self, wvMask, wRange, wvNumRegions, dC):
	# -----------------------------------------------------------------
		rois = {}

		# For stats purposes we're going to calculate the number of feasible regions as we go
		feasibleRegions = 0
		
		# Loop over the wv mask
		for w,v in wvMask.keys():
			# Update the value in the mask
			maskValue = (dC[w,v] != None)
			wvMask[w,v] = maskValue
			feasibleRegions += wvNumRegions[w,v] if maskValue else 0

		return feasibleRegions

	# -----------------------------------------------------------------	
	def generateSummaryImage(self, imageMsg, wvMask, wRange, imgHeight): 
	# -----------------------------------------------------------------
		""" Generate a |w| x |v| binary image indicating which regions are feasible """

		from PIL import Image as PImage
		#im = PImage.new('RGB', (max(wRange) + 1, imgHeight))
		im = self.bridge.imgmsg_to_cv(imageMsg, 'bgr8')

		for (w,v), value in wvMask.iteritems():
			r = 255 * (1- int(value))
			g = 255 * (int(value)) 
			b = 0
			#im.putpixel((w/2,v), (r,g,b))
			im[v, (w/2)] = (b,g,r)

		return im
			
	# --------
	# Helpers |
	# ---------------------
	def __getMessages(self):	
	# ---------------------
		""" A private function that compiles a snapshot of all of the relevant data needed to run the geometric filter. """

		# The variables we're going to try and fill
		imageMsg = None
		stereoPtCloudMsg = None
		cameraInfoMsg = None
		laserScanMsg = None
		cameraPtWorld = None
		groundNormalCamera = None

		# Make sure we have an image
		if self.imageMsg != None:
			imageMsg = self.imageMsg

			# Make sure we have a camera info message and that it's current
			ts = lambda ts: ts.secs + (ts.nsecs / 1000000000)
			imageTs = ts(imageMsg.header.stamp)
			if self.cameraInfoMsg != None and abs(imageTs - ts(self.cameraInfoMsg.header.stamp)) < config.IMG_CAMINFO_TOLERANCE_S:
				cameraInfoMsg = self.cameraInfoMsg

				# Make sure that we have a camera point					
				cameraPtWorld, groundNormalCamera = self.__getCameraOrientation(cameraInfoMsg.header.frame_id, config.BASE_FRAME)
				if cameraPtWorld != None:
					
					# If we have stereo data, include it
					if self.stereoPtCloudMsg != None and abs(imageTs - ts(self.stereoPtCloudMsg.header.stamp)) < config.IMG_STEREO_TOLERANCE_S:
						stereoPtCloudMsg = self.stereoPtCloudMsg

					# Include laser data if available
					if self.laserScanMsg != None and abs(imageTs - ts(self.laserScanMsg.header.stamp)) < config.IMG_LASER_TOLERANCE_S:
						laserScanMsg = self.laserScanMsg	

		# Return the messages we've compiled
		return (imageMsg, cameraInfoMsg, cameraPtWorld, groundNormalCamera, stereoPtCloudMsg, laserScanMsg)

	# ------------------------------------------------------------
	def __filterDepthConstraints(self, depthConstraints, maxDepth):
	# ------------------------------------------------------------
		""" Filters a list of depth constraints in the form (min, max depth).
		Depth constraints are filtered out under one of the following conditions:
		1.) The maximum depth is negative
		"""

		for i, val in depthConstraints.iteritems():
			if val[1] < 0:
				depthConstraints[i] = None
					
		return depthConstraints	

	# ------------------------------------------
	def __computeProjection(self, cameraInfoMsg):
	# ----------------------------------------- 
		""" This helper method computes the y component of the real-world ray projection for every value of v. """

		# Reshape and invert k
		k = self.__reshapeKMatrix(cameraInfoMsg.K)
		kPrime = linalg.inv(k)

		return [(((kPrime[1,1]*v) + kPrime[1,2])/((kPrime[2,1]*v) + kPrime[2,2])) for v in range(0, cameraInfoMsg.height)]						

	# ------------------------------------------------------
	def __getCameraOrientation(self, cameraFrame, baseFrame):
	# ------------------------------------------------------
		""" A private function that retrieves the position of the camera relative to some base frame.
		@param cameraFrame The tf frame of the camera.
		@param baseFrame The tf frame of the base.	

		@return a 2-tuple -- a numpy array indicating the camera's origin (x,y,z,1) and a numpy array indicating the normal vector of the ground plane (x,y,z,1)
		"""
		rospy.loginfo(cameraFrame)

		try:
			# Make sure our two frames exist
			if self.transformListener.frameExists(cameraFrame) and self.transformListener.frameExists(baseFrame):
				# Get the transformation from camera -> base frame (cb) 
				latestCommonTime = self.transformListener.getLatestCommonTime(cameraFrame, baseFrame)
				
				# Compute the camera's origin point in the base frame	
				cameraOriginBase, cbRotation = self.transformListener.lookupTransform(baseFrame, cameraFrame, latestCommonTime)

				# Compute the ground plane's normal vector in the camera's frame 
				bcTranslation, bcRotation = self.transformListener.lookupTransform(cameraFrame, baseFrame, latestCommonTime)
				bcRotationMatrix = self.transformListener.fromTranslationRotation((0,0,0), bcRotation)
				groundNormalCamera = dot(bcRotationMatrix, self.groundNormalWorld) 

				# -------------------------------------------------------------------------------------------------------------------
				# TODO: Should probably check here to make sure that groundNormalCamera[0] (x) ~ 0, otherwise the camera has 'rolled'
				# -------------------------------------------------------------------------------------------------------------------

				return (cameraOriginBase, groundNormalCamera)
			
			# Otherwise return None indicating we can't get the camera point
			else:
				return (None, None)

		# If anything goes wrong return None indicating we can't get the camera point
		except tf.Exception as e:
			return (None, None)

	# --------------------------------------------------------
	def __windowSizeRange(self, minSize, maxSize, scaleFactor):
	# --------------------------------------------------------
		""" Generates window sizes given a min/max size and a scale factor.
		Values are generated according to the formula size(i) = minSize(scaleFactor)^(i-1). 
		"""
		val = minSize
		yield val 
		while val < maxSize:
			val *= scaleFactor
			if val <= maxSize: yield int(math.floor(val))

	# ------------------------------------------------------------------------
	def __generateWVMask(self, windowSizes, imgWidth, imgHeight, wvNumRegions): 
	# ------------------------------------------------------------------------
		""" Generates a mask for all v,w combinations """

		# Initializations
		total = 0
		wvMask = {}

		for w in windowSizes:
			for v in range(int(w/2), imgHeight - int(w/2)):
				wvMask[w,v] = False			
		
				# The number of regions of interest this pairing represents	
				numRegions = imgWidth - w
				wvNumRegions[w,v] = numRegions
				total += numRegions

		################################################################
		rospy.loginfo('Total number of regions of interest: %i', total)
		self.totalROI = total
		self.numRegions = numRegions
		################################################################

		return wvMask

	# ---------------------------
	def __reshapeKMatrix(self, k):
	# ---------------------------
		""" Simple helper function to reshape a flat K matrix (presumably from a CameraInfo message) into a 3x3 numpy array. """
		if len(k) == 9:
			return array([k[0:3], k[3:6], k[6:9]])
		else: 
			raise Exception('Unable to reshape K matrix: %s' % str(k))

	# ----------------------------------------------------------------------
	def __calculateMinWindowSize(self, k, minAllowedSizeM, maxAllowedDepthM):
	# ----------------------------------------------------------------------
		""" Calculate the minimum object window size (in pixels). 
		This calculation is achieved by projecting two real world points:
		(0, 0, maxAllowedDepth) and (minAllowedSizeM, 0, maxAllowedDepth) from 3D 
		to 2D and then calculating the distance between them on the x axis.
		"""

		# Reshape and invert k
		k = self.__reshapeKMatrix(k)

		# Create our two points
		pt3A = array([[0], [0], [maxAllowedDepthM]]) 
		pt3B = array([[minAllowedSizeM], [0], [maxAllowedDepthM]]) 

		# Project the points into image space
		pt2A = dot(k, pt3A)
		pt2B = dot(k, pt3B)
	
		# Calculate the distance 
		return int(abs(pt2A[0] - pt2B[0]))

	# -----------------------------------
	def __rangeIntersection(self, r1, r2):
	# -----------------------------------
		""" Computer the intersection of two one dimensional ranges. """
		if r1 and r2:
			r = (max(r1[0], r2[0]), min(r1[1], r2[1]))
			return r if r[1] > r[0] else None 
		else: return None

	# ----------
	# Callbacks |
	# ---------------------------------
	def __cameraInfoCallback(self, msg):
	# ---------------------------------
		""" Callback for the camer information topic subscription """

		# Store the most recent camera info message we've received
		self.cameraInfoMsg = msg

	# ----------------------------
	def __imageCallback(self, msg):
	# ----------------------------
		""" Callback for the image topic subscription """		

		# Store the most recent image we've received
		self.imageMsg = msg

	# ------------------------------------
	def __stereoPtCloudCallback(self, msg):
	# ------------------------------------
		""" Callback for the stereo point cloud subscription """

		# Store the most recent stereo point cloud we've received
		self.stereoPtCloudMsg = msg

	# --------------------------------
	def __laserScanCallback(self, msg):
	# --------------------------------
		""" Callback for the laser scan topic subscription """

		# Store the most recent laser data (laser scan) that we've received
		self.laserScanMsg = msg

# -----------------------
if __name__ == '__main__':
# -----------------------
	controller = GeometricImageFilter()
	controller.run()
